
6. Fancy indexing and index tricks
===============================

NumPy offers more indexing facilities than regular Python sequences. In
addition to indexing by integers and slices, as we saw before, arrays
can be indexed by arrays of integers and arrays of booleans.

6.1 Indexing with Arrays of Indices
-------------------------------

>>> a = np.arange(12)**2                       # the first 12 square numbers
>>> i = np.array( [ 1,1,3,8,5 ] )              # an array of indices
>>> a[i]                                       # the elements of a at the positions i
array([ 1,  1,  9, 64, 25])
>>>
>>> j = np.array( [ [ 3, 4], [ 9, 7 ] ] )      # a bidimensional array of indices
>>> a[j]                                       # the same shape as j
array([[ 9, 16],
[81, 49]])

[demo]

import numpy as np
a = np.arange(12)**2                       # the first 12 square numbers
i = np.array( [ 1,1,3,8,5 ] )              # an array of indices
print(a[i])                                       # the elements of a at the positions i
j = np.array( [ [ 3, 4], [ 9, 7 ] ] )      # a bidimensional array of indices
print(a[j])

[/demo]


When the indexed array ``a`` is multidimensional, a single array of
indices refers to the first dimension of ``a``. The following example
shows this behavior by converting an image of labels into a color image
using a palette.

>>> palette = np.array( [ [0,0,0],                # black
...                       [255,0,0],              # red
...                       [0,255,0],              # green
...                       [0,0,255],              # blue
...                       [255,255,255] ] )       # white
>>> image = np.array( [ [ 0, 1, 2, 0 ],           # each value corresponds to a color in the palette
...                     [ 0, 3, 4, 0 ]  ] )
>>> palette[image]                            # the (2,4,3) color image
array([[[  0,   0,   0],
[255,   0,   0],
[  0, 255,   0],
[  0,   0,   0]],
[[  0,   0,   0],
[  0,   0, 255],
[255, 255, 255],
[  0,   0,   0]]])

[demo]

import numpy as np
palette = np.array( [ [0,0,0],                # black
                      [255,0,0],              # red
                      [0,255,0],              # green
                      [0,0,255],              # blue
                      [255,255,255] ] )       # white

image = np.array( [ [ 0, 1, 2, 0 ],           # each value corresponds to a color in the palette
                    [ 0, 3, 4, 0 ]  ] )

print(palette[image])

[/demo]


We can also give indexes for more than one dimension. The arrays of
indices for each dimension must have the same shape.

>>> a = np.arange(12).reshape(3,4)
>>> a
array([[ 0,  1,  2,  3],
[ 4,  5,  6,  7],
[ 8,  9, 10, 11]])
>>> i = np.array( [ [0,1],                        # indices for the first dim of a
...                 [1,2] ] )
>>> j = np.array( [ [2,1],                        # indices for the second dim
...                 [3,3] ] )
>>>
>>> a[i,j]                                     # i and j must have equal shape
array([[ 2,  5],
[ 7, 11]])
>>>
>>> a[i,2]
array([[ 2,  6],
[ 6, 10]])
>>>
>>> a[:,j]                                     # i.e., a[ : , j]
array([[[ 2,  1],
[ 3,  3]],
[[ 6,  5],
[ 7,  7]],
[[10,  9],
[11, 11]]])

[demo]

import numpy as np
a = np.arange(12).reshape(3,4)
print(a)
i = np.array( [ [0,1],                        # indices for the first dim of a
                [1,2] ] )

j = np.array( [ [2,1],                        # indices for the second dim
               [3,3] ] )

print(a[i,j])
print(a[i,2])
print(a[:,j])
l = [i,j]
print(a[l])

[/demo]


Naturally, we can put ``i`` and ``j`` in a sequence (say a list) and
then do the indexing with the list.

>>> l = [i,j]
>>> a[l]                                       # equivalent to a[i,j]
array([[ 2,  5],
[ 7, 11]])

However, we can not do this by putting ``i`` and ``j`` into an array,
because this array will be interpreted as indexing the first dimension
of a.

>>> s = np.array( [i,j] )
>>> a[s]                                       # not what we want
Traceback (most recent call last):
File "<stdin>", line 1, in ?
IndexError: index (3) out of range (0<=index<=2) in dimension 0
>>>
>>> a[tuple(s)]                                # same as a[i,j]
array([[ 2,  5],
[ 7, 11]])

[demo]

import numpy as np
s = np.array( [i,j] )
print(s)

[/demo]


Another common use of indexing with arrays is the search of the maximum
value of time-dependent series :

>>> time = np.linspace(20, 145, 5)                 # time scale
>>> data = np.sin(np.arange(20)).reshape(5,4)      # 4 time-dependent series
>>> time
array([  20.  ,   51.25,   82.5 ,  113.75,  145.  ])
>>> data
array([[ 0.        ,  0.84147098,  0.90929743,  0.14112001],
[-0.7568025 , -0.95892427, -0.2794155 ,  0.6569866 ],
[ 0.98935825,  0.41211849, -0.54402111, -0.99999021],
[-0.53657292,  0.42016704,  0.99060736,  0.65028784],
[-0.28790332, -0.96139749, -0.75098725,  0.14987721]])
>>>
>>> ind = data.argmax(axis=0)                   # index of the maxima for each series
>>> ind
array([2, 0, 3, 1])
>>>
>>> time_max = time[ ind]                       # times corresponding to the maxima
>>>
>>> data_max = data[ind, xrange(data.shape[1])] # => data[ind[0],0], data[ind[1],1]...
>>>
>>> time_max
array([  82.5 ,   20.  ,  113.75,   51.25])
>>> data_max
array([ 0.98935825,  0.84147098,  0.99060736,  0.6569866 ])
>>>
>>> np.all(data_max == data.max(axis=0))
True

[demo]

import numpy as np
time = np.linspace(20, 145, 5)                 # time scale
data = np.sin(np.arange(20)).reshape(5,4)      # 4 time-dependent series
print(time)
print(data)
ind = data.argmax(axis=0)                   # index of the maxima for each series
print(ind)
time_max = time[ ind]                       # times corresponding to the maxima
data_max = data[ind, xrange(data.shape[1])] # => data[ind[0],0], data[ind[1],1]...
print(time_max)
print(data_max)
print(np.all(data_max == data.max(axis=0)))

[/demo]


You can also use indexing with arrays as a target to assign to:

>>> a = np.arange(5)
>>> a
array([0, 1, 2, 3, 4])
>>> a[[1,3,4]] = 0
>>> a
array([0, 0, 2, 0, 0])

[demo]

import numpy as np
a = np.arange(5)
print(a)
a[[1,3,4]] = 0
print(a)

[/demo]


However, when the list of indices contains repetitions, the assignment
is done several times, leaving behind the last value:

>>> a = np.arange(5)
>>> a[[0,0,2]]=[1,2,3]
>>> a
array([2, 1, 3, 3, 4])

[demo]

import numpy as np
a = np.arange(5)
a[[0,0,2]]=[1,2,3]
print(a)

[/demo]


This is reasonable enough, but watch out if you want to use Python's
``+=`` construct, as it may not do what you expect:

>>> a = np.arange(5)
>>> a[[0,0,2]]+=1
>>> a
array([1, 1, 3, 3, 4])

[demo]

import numpy as np
a = np.arange(5)
a[[0,0,2]]+=1
print(a)

[/demo]

Even though 0 occurs twice in the list of indices, the 0th element is
only incremented once. This is because Python requires "a+=1" to be
equivalent to "a=a+1".

6.2 Indexing with Boolean Arrays
----------------------------

When we index arrays with arrays of (integer) indices we are providing
the list of indices to pick. With boolean indices the approach is
different; we explicitly choose which items in the array we want and
which ones we don't.

The most natural way one can think of for boolean indexing is to use
boolean arrays that have *the same shape* as the original array:

>>> a = np.arange(12).reshape(3,4)
>>> b = a > 4
>>> b                                          # b is a boolean with a's shape
array([[False, False, False, False],
[False,  True,  True,  True],
[ True,  True,  True,  True]], dtype=bool)
>>> a[b]                                       # 1d array with the selected elements
array([ 5,  6,  7,  8,  9, 10, 11])

This property can be very useful in assignments:

>>> a[b] = 0                                   # All elements of 'a' higher than 4 become 0
>>> a
array([[0, 1, 2, 3],
[4, 0, 0, 0],
[0, 0, 0, 0]])

[demo]

import numpy as np
a = np.arange(12).reshape(3,4)
b = a > 4
print(b)
print(a[b])
a[b] = 0
print(a)

[/demo]


You can look at the following
example to see
how to use boolean indexing to generate an image of the `Mandelbrot
set <http://en.wikipedia.org/wiki/Mandelbrot_set>`__:

.. plot::

>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> def mandelbrot( h,w, maxit=20 ):
...     """Returns an image of the Mandelbrot fractal of size (h,w)."""
...     y,x = np.ogrid[ -1.4:1.4:h*1j, -2:0.8:w*1j ]
...     c = x+y*1j
...     z = c
...     divtime = maxit + np.zeros(z.shape, dtype=int)
...
...     for i in range(maxit):
...         z = z**2 + c
...         diverge = z*np.conj(z) > 2**2            # who is diverging
...         div_now = diverge & (divtime==maxit)  # who is diverging now
...         divtime[div_now] = i                  # note when
...         z[diverge] = 2                        # avoid diverging too much
...
...     return divtime
>>> plt.imshow(mandelbrot(400,400))
>>> plt.show()

[demo]

import numpy as np
import matplotlib.pyplot as plt
def mandelbrot( h,w, maxit=20 ):
    """Returns an image of the Mandelbrot fractal of size (h,w)."""
    y,x = np.ogrid[ -1.4:1.4:h*1j, -2:0.8:w*1j ]
    c = x+y*1j
    z = c
    divtime = maxit + np.zeros(z.shape, dtype=int)
    for i in range(maxit):
        z = z**2 + c
        diverge = z*np.conj(z) > 2**2            # who is diverging
        div_now = diverge & (divtime==maxit)  # who is diverging now
        divtime[div_now] = i                  # note when
        z[diverge] = 2                        # avoid diverging too much
    return divtime

plt.imshow(mandelbrot(400,400))
plt.show()

[/demo]


The second way of indexing with booleans is more similar to integer
indexing; for each dimension of the array we give a 1D boolean array
selecting the slices we want.

>>> a = np.arange(12).reshape(3,4)
>>> b1 = np.array([False,True,True])             # first dim selection
>>> b2 = np.array([True,False,True,False])       # second dim selection
>>>
>>> a[b1,:]                                   # selecting rows
array([[ 4,  5,  6,  7],
[ 8,  9, 10, 11]])
>>>
>>> a[b1]                                     # same thing
array([[ 4,  5,  6,  7],
[ 8,  9, 10, 11]])
>>>
>>> a[:,b2]                                   # selecting columns
array([[ 0,  2],
[ 4,  6],
[ 8, 10]])
>>>
>>> a[b1,b2]                                  # a weird thing to do
array([ 4, 10])


[demo]

import numpy as np
a = np.arange(12).reshape(3,4)
b1 = np.array([False,True,True])             # first dim selection
b2 = np.array([True,False,True,False])       # second dim selection
print(a[b1,:] )                                  # selecting rows
print(a[b1])                                    # same thing
print(a[:,b2])                                 # selecting columns
print(a[b1,b2])

[/demo]

Note that the length of the 1D boolean array must coincide with the
length of the dimension (or axis) you want to slice. In the previous
example, ``b1`` is a 1-rank array with length 3 (the number of *rows* in
``a``), and ``b2`` (of length 4) is suitable to index the 2nd rank
(columns) of ``a``.

6.3 The ix_() function
-------------------

The `ix_` function can be used to combine different vectors so as to
obtain the result for each n-uplet. For example, if you want to compute
all the a+b\*c for all the triplets taken from each of the vectors a, b
and c:

>>> a = np.array([2,3,4,5])
>>> b = np.array([8,5,4])
>>> c = np.array([5,4,6,8,3])
>>> ax,bx,cx = np.ix_(a,b,c)
>>> ax
array([[[2]],
[[3]],
[[4]],
[[5]]])
>>> bx
array([[[8],
[5],
[4]]])
>>> cx
array([[[5, 4, 6, 8, 3]]])
>>> ax.shape, bx.shape, cx.shape
((4, 1, 1), (1, 3, 1), (1, 1, 5))
>>> result = ax+bx*cx
>>> result
array([[[42, 34, 50, 66, 26],
[27, 22, 32, 42, 17],
[22, 18, 26, 34, 14]],
[[43, 35, 51, 67, 27],
[28, 23, 33, 43, 18],
[23, 19, 27, 35, 15]],
[[44, 36, 52, 68, 28],
[29, 24, 34, 44, 19],
[24, 20, 28, 36, 16]],
[[45, 37, 53, 69, 29],
[30, 25, 35, 45, 20],
[25, 21, 29, 37, 17]]])
>>> result[3,2,4]
17
>>> a[3]+b[2]*c[4]
17

You could also implement the reduce as follows:

>>> def ufunc_reduce(ufct, *vectors):
...    vs = np.ix_(*vectors)
...    r = ufct.identity
...    for v in vs:
...        r = ufct(r,v)
...    return r

and then use it as:

>>> ufunc_reduce(np.add,a,b,c)
array([[[15, 14, 16, 18, 13],
[12, 11, 13, 15, 10],
[11, 10, 12, 14,  9]],
[[16, 15, 17, 19, 14],
[13, 12, 14, 16, 11],
[12, 11, 13, 15, 10]],
[[17, 16, 18, 20, 15],
[14, 13, 15, 17, 12],
[13, 12, 14, 16, 11]],
[[18, 17, 19, 21, 16],
[15, 14, 16, 18, 13],
[14, 13, 15, 17, 12]]])


[demo]

import numpy as np
a = np.array([2,3,4,5])
b = np.array([8,5,4])
c = np.array([5,4,6,8,3])
ax,bx,cx = np.ix_(a,b,c)
print(ax)
print(bx)
print(cx)
print(ax.shape, bx.shape, cx.shape)
result = ax+bx*cx
print(result)
print(result[3,2,4])
print(a[3]+b[2]*c[4])

def ufunc_reduce(ufct, *vectors):
   vs = np.ix_(*vectors)
   r = ufct.identity
   for v in vs:
       r = ufct(r,v)
   return r

print(ufunc_reduce(np.add,a,b,c))

[/demo]

The advantage of this version of reduce compared to the normal
ufunc.reduce is that it makes use of the `Broadcasting
Rules <Tentative_NumPy_Tutorial.html#head-c43f3f81719d84f09ae2b33a22eaf50b26333db8>`__
in order to avoid creating an argument array the size of the output
times the number of vectors.

6.4 Indexing with strings
---------------------

See `RecordArrays <RecordArrays.html>`__.